package pers.qianyu.month_202102.date_20210209;

import org.junit.*;

/**
 * 992. K 个不同整数的子数组
 * https://leetcode-cn.com/problems/subarrays-with-k-different-integers/
 * - 参考题解：https://leetcode-cn.com/problems/subarrays-with-k-different-integers/solution/k-ge-bu-tong-zheng-shu-de-zi-shu-zu-by-l-ud34/
 *
 * @author mizzle rain
 * @date 2021-02-09 22:11
 */
public class SubarraysWithKDistinct {
    public int subarraysWithKDistinct(int[] nums, int k) {
        return atMostWithKDistinct(nums, k) - atMostWithKDistinct(nums, k - 1);
    }

    /**
     * 最多包含 K 个不同整数的子区间个数
     */
    private int atMostWithKDistinct(int[] nums, int k) {
        int n = nums.length;
        int left = 0, right = 0;
        int cnt = 0;
        int[] map = new int[n + 1];
        int res = 0;
        while (right < n) {
            if (map[nums[right]] == 0) cnt++;
            map[nums[right]]++;
            while (cnt > k) {
                map[nums[left]]--;
                if (map[nums[left]] == 0) cnt--;
                left++;
            }
            right++;
            res += right - left;
        }
        return res;
    }

    @Test
    public void test1() {
        int result = new SubarraysWithKDistinct()
                .subarraysWithKDistinct(new int[]{1, 2, 1, 2, 3}, 2);
        Assert.assertEquals(7, result);
    }
}
